#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 05:57:30
# @LastEditTime: 2025-03-14 05:57:35
# @Description: MySQL数据库SQL文件导入工具
###

import os
import sys
import pymysql
import subprocess
import configparser
from tqdm import tqdm
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mysql_import.log'),
        logging.StreamHandler()
    ]
)

def load_config():
    """从配置文件加载数据库连接信息"""
    config_file = os.path.join(os.path.dirname(__file__), 'config.ini')
    if not os.path.exists(config_file):
        logging.error(f"错误：配置文件不存在：{config_file}")
        sys.exit(1)
    
    config = configparser.ConfigParser()
    config.read(config_file)
    
    if 'restore' not in config:
        logging.error("错误：配置文件中缺少 [restore] 部分")
        sys.exit(1)
        
    return dict(config['restore'])

def create_database(db_config, db_name):
    """创建数据库（如果不存在）"""
    try:
        connection = pymysql.connect(**db_config)
        cursor = connection.cursor()
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;")
        connection.commit()
        logging.info(f"数据库 {db_name} 准备就绪")
        return True
    except pymysql.Error as err:
        logging.error(f"创建数据库失败: {err}")
        return False
    finally:
        cursor.close()
        connection.close()

def import_sql_file(db_config, db_name, file_path):
    """导入SQL文件到指定数据库"""
    try:
        file_size = os.path.getsize(file_path)
        logging.info(f"开始导入 {file_path} (大小: {file_size/1024/1024:.2f}MB) 到数据库 {db_name}")
        
        mysql_cmd = (
            f"mysql -u {db_config['user']} -p{db_config['password']} "
            f"-h {db_config['host']} -P {db_config['port']} "
            f"--default-character-set=utf8mb4 {db_name} < {file_path}"
        )
        
        process = subprocess.run(
            mysql_cmd, 
            shell=True, 
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE
        )
        
        if process.returncode != 0:
            logging.error(f"导入失败：{process.stderr.decode()}")
            return False
            
        logging.info(f"成功导入 {file_path} 到数据库 {db_name}")
        return True
    except Exception as e:
        logging.error(f"导入过程发生错误：{e}")
        return False

def main():
    """主函数"""
    try:
        # 加载配置
        db_config = load_config()
        
        # 获取当前目录下的所有SQL文件
        current_dir = os.getcwd()
        sql_files = [f for f in os.listdir(current_dir) if f.endswith('.sql')]
        
        if not sql_files:
            logging.error("当前目录下没有找到.sql文件")
            sys.exit(1)
        
        logging.info(f"找到 {len(sql_files)} 个SQL文件待导入")
        
        # 使用进度条显示导入进度
        with tqdm(sql_files, desc="导入进度") as pbar:
            for sql_file in pbar:
                db_name = os.path.splitext(sql_file)[0]  # 去掉.sql后缀获取数据库名
                file_path = os.path.join(current_dir, sql_file)
                
                pbar.set_description(f"正在处理 {db_name}")
                
                # 创建数据库
                if not create_database(db_config, db_name):
                    continue
                
                # 导入SQL文件
                import_sql_file(db_config, db_name, file_path)
                
        logging.info("所有导入操作已完成")
        
    except KeyboardInterrupt:
        logging.warning("操作被用户中断")
        sys.exit(1)
    except Exception as e:
        logging.error(f"发生错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()